import json,uuid,os,aiofiles,magic
from wand.image import Image
from cnocr import CnOcr
from cnstd import CnStd
import subprocess
import sys, fitz
import datetime
import docx
import pdfplumber
from exception import *

class OcrHelper:

    settings = {
        'MEDIA_ROOT': './temp/'
    }

    # 对PDF进行截图
    def pyMuPDF_fitz(self, pdfPath, imagePath):
        startTime_pdf2img = datetime.datetime.now()#开始时间
        
        resPath = []
        pdfDoc = fitz.open(pdfPath)
        for pg in range(pdfDoc.pageCount):
            page = pdfDoc[pg]
            rotate = int(0)
            # 每个尺寸的缩放系数为1.3，这将为我们生成分辨率提高2.6的图像。
            # 此处若是不做设置，默认图片大小为：792X612, dpi=96
            zoom_x = 1.33333333 #(1.33333333-->1056x816)   (2-->1584x1224)
            zoom_y = 1.33333333
            mat = fitz.Matrix(zoom_x, zoom_y).preRotate(rotate)
            pix = page.getPixmap(matrix=mat, alpha=False)
            
            if not os.path.exists(imagePath):#判断存放图片的文件夹是否存在
                os.makedirs(imagePath) # 若图片文件夹不存在就创建
            
            filePath = imagePath+'/'+'images_%s.png' % pg
            pix.writePNG(filePath)#将图片写入指定的文件夹内
            resPath.append(filePath)
            
        endTime_pdf2img = datetime.datetime.now()#结束时间
        # print('pdf2img时间=',(endTime_pdf2img - startTime_pdf2img).seconds)
        return resPath
    
    def pdf2image(self, file_path):
        with(Image(filename=file_path, resolution=120)) as source: 
            images = source.sequence
            pages = len(images)
            for i in range(pages):
                n = i + 1 
                newfilename = file_path[:-4] + str(n) + '.jpeg'
                Image(images[i]).save(filename=newfilename)

    # 对PDF截图再进行OCR识别
    def pdfWithImage(self, file_path):
        # 先截图
        image_paths = self.pyMuPDF_fitz(file_path, self.settings["MEDIA_ROOT"])
        results = []
        for image_path in image_paths:
            # print(image_path)
            # 再OCR识别
            results = results + self.image_recognition(image_path)
            os.remove(image_path)

        return results
    
    # 直接使用pdfplumber去识别
    def pdfplumber(self, file_path):
        results = []
        pdf = pdfplumber.open(file_path)
        for page in pdf.pages:
            content = page.extract_text().replace(' ', '').replace('space', '').replace('\n', '')
            results.append(content)
            tables = page.extract_tables()
            for table in tables:
                for row in table:
                    results.append(str(row))
        pdf.close()
        return results

    # doc转docx
    def doc2docx(self, docPath, docxPath):
        cmd = 'libreoffice --headless --convert-to docx'.split() + [docPath] + ['--outdir'] + [docxPath]
        p = subprocess.Popen(cmd, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
        p.wait(timeout=60)
        stdout, stderr = p.communicate()
        # print(stdout)
        if stderr:
            raise subprocess.SubprocessError(stderr)
    
    # doc转jpeg/png
    def doc2other(self, docPath, docxPath, type='docx'):
        cmd = 'libreoffice --headless --convert-to {}'.format(type).split() + [docPath] + ['--outdir'] + [docxPath]
        print(cmd)
        subprocess.call(cmd)
        # subprocess.Popen(cmd, shell=True)
            
    # docx转文字
    def docx2text(self, file_path):
        document = docx.Document(file_path)
        # print(file_path)
        #获取所有段落
        all_paragraphs = document.paragraphs
        all_texts = []
        for paragraph in all_paragraphs:
            all_texts.append(paragraph.text)
        
        tables = document.tables
        for table in tables:
            table_data = self.docx_table(table)
            all_texts.append(str(table_data['head']))
            all_texts.append(str(table_data['mat']))
        return all_texts
    
    # 读取docx表格中的数据
    def docx_table(self, table):
        mat = [] #存储行
        mat_rows = [] #临时存储行
        head = [] #存储表头
        columns = table._column_count #列的长度 
        rows = len(table.rows) #行的长度

        for c in range(0,columns):
            cell = table.cell(0,c)
            txt = cell.text.replace('\n', '')
            head.append(txt)

        #将表头追加到 head 列表中

        #先将所有的行取出来
        for r in range(1,rows):
            mat_rows.append(table.row_cells(r))

        #将取出来的行进行拆解
        for rows_datas in mat_rows:
            row = []
            for i in rows_datas:
                txt = i.text if i.text != '' else ''
                row.append(txt)
            mat.append(row)
        
        return {"head": head, "mat": mat, "columns": columns, "rows": rows}

    # CNOCR识别
    def image_recognition(self, file):
        
        #CNOCR是将 文字检测模型和文字识别模型分开，我下面注释的部分是，通过调用不同的模型来看识别效果和识别速度，大家可以根据自己的需要来调整不同的模型，来应用到自己的场景中。
        
        # ocr = CnOcr()
        ocr = CnOcr(det_model_name='naive_det')
        res = ocr.ocr(file)
        # print(res)
        ret_data = [x['text'] for x in res]
        if len(ret_data) == 0 or ret_data[0] == "":
            ocr = CnOcr(rec_model_name='ch_PP-OCRv3')
            res = ocr.ocr(file)
            # print(res)
            ret_data = [x['text'] for x in res]

        # ocr = CnOcr(det_model_name='db_resnet34', rec_model_name='ch_PP-OCRv3')    #使用resnet34模型检测文字，ch_PP-OCRV3来识别文字
        # ocr = CnOcr(det_model_name='db_shufflenet_v2_small', rec_model_name='en_PP-OCRv3')
        # ocr = CnOcr(det_model_name='ch_PP-OCRv3_det', rec_model_name='densenet_lite_136-fc',)
        # ocr = CnOcr(det_model_name='en_PP-OCRv3_det', rec_model_name='en_number_mobile_v2.0',det_more_configs={'resized_shape':1200})
        # ocr = CnOcr(rec_model_name='ch_PP-OCRv3')

        # res = ocr.ocr_for_single_line(file) 
        # print("Predicted Chars:", res)
        # print('================================')
        # res = ocr.ocr('mozhang8.png')
        # res = sorted(res,key=lambda x:x['position'][0][0])
        # print()
        # print([x['text'] for x in res])
        # res = ocr.ocr('itemsaaa.png')
        # print(res)
        # astd = CnStd('en_PP-OCRv3_det')
        # glist = []
        # # for x in astd.detect('itemsaaa.png',resized_shape=1200,).get('detected_texts'):
        # for x in sorted(ocr.det_model.detect(file,resized_shape=1200,).get('detected_texts'),key=lambda m:m['box'][0][0]):
        #     cropped_img = x['cropped_img']
        #     ocr_res = ocr.ocr_for_single_line(cropped_img)
        #     # ocr.ocr_for_single_lines()
        #     result = {'box':x['box']}
        #     result.update(ocr_res)
        #     glist.append(result)
        #     print('ocr result: %s' % str(ocr_res))
        
        # ret_data = [x['text'] for x in glist]
        # ret_data.reverse()
        # print(ret_data)
        return list(filter(None, ret_data))
        
        # images = sorted(ocr.det_model.detect('itemsaaa.png').get('detected_texts'),key=lambda m:m['box'][0][0])
        
        #以下的代码是将需要识别的图片进行 resize，这个参数会对识别结果和识别速度有影响，分辨率越大，识别速度越慢。当然也不是越大就是越精准。
        # images = sorted(ocr.det_model.detect(file,resized_shape=1200,).get('detected_texts'),key=lambda m:m['box'][0][0])
        # imgs = [ x['cropped_img'] for x in images]
        # ocr_res = ocr.ocr_for_single_lines(imgs)
        # print([x['text'] for x in ocr_res])
        # print('================================')
    
    def handle(self, file_path):

        ret_data = {
            'file_path': file_path,
            'result': None
        }
        # print(file_path)
        file_type = file_path[file_path.rfind('.')+1:].lower()

        try:

            # 图片
            if file_type in ['png', 'jpg', 'jpeg', 'bmp']:
                ret_data['result'] = self.image_recognition(file_path)

            # text
            elif file_type == 'txt':
                f = open(file_path, encoding="utf-8")
                ret_data['result'] = f.readlines()
                f.close()

            # WPS PDF
            elif file_type == 'pdf':
                try:
                    # 直接识别PDF
                    # print('直接识别pdf')
                    results = self.pdfplumber(file_path)
                    ret_data['result'] = results
                except Exception as err:
                    print(err)
                    # print('截图识别pdf')
                    # 通过截图识别PDF
                    try:
                        results = self.pdfWithImage(file_path)
                        ret_data['result'] = results
                    except Exception as e:
                        print(e)
                        # resp_data['code'] = 5
                        # resp_data['message'] = "pdfWithImage识别错误：%s"%e
                        raise OcrError("pdfWithImage识别错误：%s"%e)
                
            # WPS DOC
            elif file_type == 'doc':
                # print('doc')
                to_type = 'docx'
                self.doc2other(file_path, self.settings["MEDIA_ROOT"], to_type)
                new_path = file_path[:-3] + to_type
                # ret_data['result'] = self.image_recognition(new_path)
                print(new_path)
                ret_data['result'] = self.docx2text(new_path)
                os.remove(new_path)

            # WPS DOCX
            # elif file_type[1] == 'vnd.openxmlformats-officedocument.wordprocessingml.document':
            elif file_type == 'docx':
                ret_data['result'] = self.docx2text(file_path)

            # # WPS XLS 
            # elif file_type[1] == 'vnd.ms-excel':
            #     print('xls')

            # # WPS XLSX
            # elif file_type[1] == 'vnd.openxmlformats-officedocument.spreadsheetml.sheet':
            #     print('xlsx')

            else:
                # resp_data['code'] = 3
                # resp_data['message'] = "不支持的文件格式：%s"%file_type
                raise TypeError("不支持的文件格式：%s"%file_type)

        except Exception as error:
            raise Exception("error：%s"%error)

        return ret_data
        